import os
import dashscope
import clingo
from datetime import datetime

# 在代码最开头明确定义 ASP 代码（重要！）
asp_code = """
#const num_scenes = 7.
scene(1..num_scenes).
function(intro_char; intro_rival_char; add_conflict; add_bonding; 
        add_obstacle_towards_major_goal; add_obstacle; add_breakthrough; 
        add_twist; level_up_obstacle; final_battle; generate_ending).

% 每个场景有一个功能
1 { scene_performs_func(S,F) : function(F) } 1 :- scene(S).

% 强制最后两个场景的结构
:- scene_performs_func(6, F), F != final_battle.  % 第6场景必须是最终战斗
:- scene_performs_func(7, F), F != generate_ending.  % 第7场景必须是结局

% 原有约束保持不变
:- scene_performs_func(S1,F), scene_performs_func(S2,F), S2 = S1 + 1.
:- scene_performs_func(S2,add_conflict), 0 { scene_performs_func(S1,intro_char) : scene(S1), S1 < S2 } 1.
:- scene_performs_func(S,add_twist), S < 4.
"""
# 在代码开头（建议在asp_code定义之后）明确定义翻译字典
# 更新后的翻译字典（包含最终战斗和结局生成）
translation_dict = {
    # 角色介绍类
    "intro_char":
        "Introduce a new important character. Should include:\n"
        "- Physical description\n"
        "- Personality traits\n"
        "- Initial motivation/goals",

    "intro_rival_char":
        "Introduce an antagonist character. Must include:\n"
        "- Source of conflict with protagonist\n"
        "- Unique abilities/resources\n"
        "- Personality flaws that drive opposition",

    # 冲突发展类
    "add_conflict":
        "Develop a major conflict sequence. Requirements:\n"
        "- Clear opposing forces\n"
        "- High emotional stakes\n"
        "- Consequences for failure",

    "add_bonding":
        "Create a character bonding scene. Elements:\n"
        "- Shared vulnerability moment\n"
        "- Exchange of meaningful gifts/information\n"
        "- Establishment of mutual trust",

    # 障碍设置类
    "add_obstacle_towards_major_goal":
        "Create a primary goal obstacle. Must:\n"
        "- Directly block main objective\n"
        "- Require multi-step solution\n"
        "- Test protagonist's core strength",

    "add_obstacle":
        "Add secondary obstacle. Should:\n"
        "- Create immediate pressure\n"
        "- Reveal character weaknesses\n"
        "- Require quick improvisation",

    "level_up_obstacle":
        "Escalate existing obstacle by:\n"
        "1. Adding time/resource constraints\n"
        "2. Introducing moral dilemmas\n"
        "3. Raising personal stakes",

    # 转折突破类
    "add_breakthrough":
        "Create a breakthrough moment with:\n"
        "- Unexpected resource acquisition\n"
        "- Key ally intervention\n"
        "- Crucial self-realization",

    "add_twist":
        "Implement plot twist that:\n"
        "- Recontextualizes prior events\n"
        "- Reveals hidden connections\n"
        "- Forces goal reassessment",

    # 新增最终阶段功能
    "final_battle":
        "Write climactic battle sequence with:\n"
        "1. Three-phase structure (setup/confrontation/resolution)\n"
        "2. Physical and emotional stakes\n"
        "3. Sacrifices from major characters\n"
        "4. Demonstration of character growth",

    "generate_ending":
        "Create story resolution that:\n"
        "- Shows consequences of key choices\n"
        "- Resolves character arcs\n"
        "- Leaves lasting thematic impression\n"
        "- Provides closure with hint of future possibilities"
}

def print_log(message, level="INFO"):
    """统一的日志输出函数"""
    timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    print(f"[{timestamp}] [{level}] {message}")


def generate_outline(asp_content=None):
    """
    生成符合叙事结构的大纲（强制包含最终战斗和结局）
    返回格式：[func1, func2, ..., final_battle, generate_ending]
    """
    print_log("开始生成故事大纲...")

    # ================= 初始化阶段 =================
    used_asp_code = None
    try:
        # 代码来源优先级：参数 > 全局变量
        if asp_content:
            print_log("使用自定义ASP代码")
            used_asp_code = asp_content
        else:
            global asp_code
            if 'asp_code' in globals() and isinstance(asp_code, str) and len(asp_code) > 100:
                print_log("使用全局ASP代码")
                used_asp_code = asp_code
            else:
                print_log("ASP代码无效：类型错误或长度不足", "ERROR")
                return []

        # 基础验证
        if not used_asp_code.strip().startswith("#const"):
            print_log("ASP代码缺少必要约束声明", "WARNING")

    except Exception as e:
        print_log(f"初始化失败：{str(e)}", "ERROR")
        return []

    # ================= Clingo 配置 =================
    control = clingo.Control()
    try:
        print_log("加载ASP代码...")
        control.add("base", [], used_asp_code)

        print_log("实例化基础规则...")
        control.ground([("base", [])])

        # 配置求解器参数
        control.configuration.solve.parallel_mode = 4  # 启用多线程
        control.configuration.solve.models = 0  # 查找所有解
        print_log(f"求解器配置完成：{control.configuration.solve}")

    except RuntimeError as e:
        print_log(f"Clingo初始化错误：{str(e)}", "ERROR")
        return []

    # ================= 模型求解阶段 =================
    valid_outlines = []
    try:
        print_log("开始求解过程...")
        with control.solve(yield_=True) as handle:
            for model_idx, model in enumerate(handle, 1):
                print_log(f"处理模型 #{model_idx}")

                # 解析原子数据
                scene_funcs = []
                for atom in model.symbols(atoms=True):
                    if atom.name == "scene_performs_func":
                        try:
                            scene = atom.arguments[0].number
                            func = atom.arguments[1].name
                            scene_funcs.append((scene, func))
                        except AttributeError:
                            print_log(f"无效原子结构：{atom}", "WARNING")

                # 按场景排序
                scene_funcs.sort(key=lambda x: x[0])
                current_outline = [func for _, func in scene_funcs]

                # 结构验证
                if len(current_outline) >= 7:
                    last_two = current_outline[-2:]
                    if last_two == ['final_battle', 'generate_ending']:
                        print_log(f"模型 #{model_idx} 通过最终结构验证")
                        valid_outlines.append(current_outline)
                    else:
                        print_log(f"模型 #{model_idx} 结构无效：最后两个场景为 {last_two}", "WARNING")
                else:
                    print_log(f"模型 #{model_idx} 场景数量不足：{len(current_outline)}", "WARNING")

                # 安全限制
                if model_idx >= 10:
                    print_log("达到最大模型检查限制（10个）", "WARNING")
                    break

    except Exception as e:
        print_log(f"求解过程异常：{str(e)}", "ERROR")
        return []

    # ================= 结果处理 =================
    if valid_outlines:
        print_log(f"找到 {len(valid_outlines)} 个有效大纲")
        # 优先选择最长的有效大纲
        final_outline = max(valid_outlines, key=len)
        print_log("最终采用大纲：" + " -> ".join(final_outline))
        return final_outline
    else:
        print_log("未找到符合约束的大纲", "ERROR")
        # 应急生成逻辑
        return [
            'intro_char', 'add_conflict', 'add_obstacle',
            'add_twist', 'level_up_obstacle',
            'final_battle', 'generate_ending'
        ]
def translate_outline(outline):
    """将大纲中的功能翻译为写作指令"""
    print_log("开始翻译大纲为指令...")
    try:
        # 添加字典存在性检查
        if 'translation_dict' not in globals():
            print_log("错误：translation_dict未定义", "ERROR")
            return []

        translated = []
        for idx, func in enumerate(outline, 1):
            # 添加键存在性检查
            if func in translation_dict:
                print_log(f"场景{idx}: 转换功能 '{func}' => '{translation_dict[func]}'")
                translated.append(translation_dict[func])
            else:
                print_log(f"场景{idx}: 发现未知功能 '{func}'，已跳过", "WARNING")
        return translated
    except Exception as e:
        print_log(f"翻译大纲时发生异常：{str(e)}", "ERROR")
        return []

def generate_story(premise, instructions):
    """使用DashScope生成英语故事"""
    print_log("开始生成故事内容...")
    story = ""
    system_prompt = "You are a professional English story writer."

    if not instructions:
        print_log("没有可用的写作指令，跳过生成步骤", "WARNING")
        return ""

    for idx, instruction in enumerate(instructions, 1):
        print_log(f"正在处理第 {idx}/{len(instructions)} 个场景指令")
        print_log(f"当前指令: {instruction}")

        try:
            # 构建提示语
            if idx == 1:
                prompt = f"Premise: {premise}\nInstruction: {instruction}\nGenerate the first paragraph"
                print_log("生成首段提示语")
            else:
                prompt = f"Premise: {premise}\nInstruction: {instruction}\nCurrent story: {story[-500:]}\nContinue the next paragraph"  # 只保留最近500字符避免过长
                print_log("生成续写提示语（带上下文截断）")

            # 显示API调用参数（不显示敏感信息）
            print_log(f"调用API参数：model=qwen-max, prompt_length={len(prompt)}")

            # 检查API密钥
            if not os.getenv('DASHSCOPE_API_KEY'):
                print_log("未检测到DASHSCOPE_API_KEY环境变量！", "ERROR")
                return story + "\n[ERROR: MISSING API KEY]"

            # 调用API
            response = dashscope.Generation.call(
                api_key=os.getenv('DASHSCOPE_API_KEY'),
                model="qwen-plus-2025-01-25",
                messages=[
                    {'role': 'system', 'content': system_prompt},
                    {'role': 'user', 'content': prompt}
                ],
                result_format='message'
            )

            # 处理响应
            print_log(f"收到API响应，状态码：{response.status_code}")
            if response.status_code == 200:
                if hasattr(response, 'output') and hasattr(response.output, 'choices'):
                    if len(response.output.choices) > 0:
                        content = response.output.choices[0].message.get('content', '')
                        print_log(f"成功获取内容，长度：{len(content)}字符")
                        story += content + "\n\n"
                        print_log(f"当前故事总长度：{len(story)}字符")
                    else:
                        print_log("API返回空的choices列表", "WARNING")
                        story += "[CONTENT MISSING]\n\n"
                else:
                    print_log("意外的API响应结构", "ERROR")
                    print_log(f"完整响应结构：{response}")
                    story += "[RESPONSE FORMAT ERROR]\n\n"
            else:
                print_log(f"API调用失败，状态码：{response.status_code}", "ERROR")
                print_log(f"错误详情：{response}")
                story += f"[API ERROR {response.status_code}]\n\n"

        except Exception as e:
            print_log(f"生成场景{idx}时发生异常：{str(e)}", "ERROR")
            story += f"[GENERATION ERROR: {str(e)}]\n\n"

        print_log("-" * 50)  # 分割线

    return story


if __name__ == "__main__":
    # 示例前提（英语）
    premise = """In a faraway kingdom, there lived a beautiful young girl named Ella. After her father's sudden death,
    she was forced to work as a servant for her wicked stepmother and two jealous stepsisters."""

    print_log("故事生成流程启动")
    print_log(f"初始前提：\n{premise}")

    # 生成大纲
    outline = generate_outline()
    print_log(f"原始大纲结构：{outline}")

    # 翻译大纲
    instructions = translate_outline(outline)
    print_log(f"最终写作指令：")
    for i, instr in enumerate(instructions, 1):
        print_log(f"指令{i}: {instr}")

    # 生成故事
    story = generate_story(premise, instructions)

    # 最终输出
    print_log("故事生成完成")
    print("\n" + "=" * 50 + " FINAL STORY " + "=" * 50)
    print(story)